import os
import random

import numpy as np
import torch
import torchvision.datasets as datasets
from PIL import Image
from torch.utils.data.dataset import Dataset
from torchvision import transforms
from .utils import cvtColor, preprocess_input, resize_image

class FacenetDataset_re(Dataset):
    def __init__(self, input_shape, lines, num_classes, mode='train'):
        self.input_shape    = input_shape
        self.lines          = lines
        self.length         = len(lines)
        self.num_classes    = num_classes
        self.mode = mode
        assert self.mode=='train' or self.mode=='test', "dataset gets a wrong setting"
        
        self.paths  = []
        self.labels = []
        '''
            trans = transforms.Compose([
            transforms.CenterCrop([120, 120]),
            transforms.RandomRotation(20),
            transforms.RandomAdjustSharpness(sharpness_factor=0.4 ,p=0.5),
            transforms.RandomGrayscale(p=0.3),
            transforms.ToTensor(),
            transforms.Normalize(mean = [0.485, 0.456, 0.406], 
                                 std = [0.229, 0.224, 0.225]),
            transforms.RandomErasing(p=0.5, scale=(0.02, 0.23), ratio=(0.3, 3.3)),
        ])
        '''
        # Imagenet数据集的均值和方差为：mean=(0.485, 0.456, 0.406)，std=(0.229, 0.224, 0.225)
        # 要先把图片 transform.ToTensor()，再进行RandomErasing()，也就是说随机擦除的过程是在tensor数据上进行的，这与上面的裁剪，翻转不一样。
        self.crop = transforms.CenterCrop((120,120))  # [120, 120]
        self.norm = transforms.Normalize(mean = [0.485, 0.456, 0.406], 
                                         std = [0.229, 0.224, 0.225])
        self.resize = transforms.Resize(size=(160,160))  # [160, 160]
        if mode == 'train':
            self.train_trans0 = transforms.RandomHorizontalFlip(p=0.5)
            self.train_trans0_5 = transforms.RandomGrayscale(p=0.2)
            self.train_trans1 = transforms.RandomRotation(20)
            self.train_trans2 = transforms.RandomAdjustSharpness(sharpness_factor=0.4 ,p=0.5)
            self.train_trans3 = transforms.RandomErasing(p=0.5, scale=(0.02, 0.23), ratio=(0.3, 3.3))
            self.train_trans4 = transforms.RandomCrop(size=(90,90))
            self.train_trans5 = transforms.RandomPerspective(distortion_scale=0.4, p=1.0)
        elif mode == 'test':
            self.trans = transforms.Compose([
                # transforms.RandomRotation(20),
            ])
        
        self.load_dataset()
        
    def __len__(self):
        return self.length
    
    # acc:95.7
#     def __getitem__(self, index):
#         image = Image.open(self.paths[index])
#         label = self.labels[index]
#         image = self.crop(image)
#         image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image = True)
#         if self.mode == 'train':
#             image = self.train_trans0(image)
#             image = self.train_trans1(image)
#             image = self.train_trans2(image)
            
#             image = preprocess_input(np.array(image, dtype='float32'))
#             image = np.transpose(image, [2, 0, 1])
#             image = torch.from_numpy(image)
#             image = self.norm(image)
            
#             image = self.train_trans3(image)
        
#         elif self.mode == 'test':
#             image = preprocess_input(np.array(image, dtype='float32'))
#             image = np.transpose(image, [2, 0, 1])
#             image = torch.from_numpy(image)
#             image = self.norm(image)
        
#         return image, label

    # acc:96.1
#     def __getitem__(self, index):
#         image = Image.open(self.paths[index])
#         label = self.labels[index]
#         image = self.crop(image)
#         image = self.resize(image)
#         # image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image = True)
#         if self.mode == 'train':
#             image = self.train_trans0(image)
#             image = self.train_trans1(image)
#             image = self.train_trans2(image)
            
#             image = preprocess_input(np.array(image, dtype='float32'))
#             image = np.transpose(image, [2, 0, 1])
#             image = torch.from_numpy(image)
#             image = self.norm(image)
            
#             image = self.train_trans3(image)
        
#         elif self.mode == 'test':
#             image = preprocess_input(np.array(image, dtype='float32'))
#             image = np.transpose(image, [2, 0, 1])
#             image = torch.from_numpy(image)
#             image = self.norm(image)
        
#         return image, label

    # acc: 96.5  acc: 97.3(with cplfw)
    def __getitem__(self, index):
        image = Image.open(self.paths[index])
        label = self.labels[index]
        image = self.crop(image)  # 160 160
        # image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image = True)
        if self.mode == 'train':
            prob = np.random.rand()
            if prob > 0.5:
                image = self.train_trans4(image)  # 90 90
            image = self.resize(image)  # 160 160
            image = self.train_trans0_5(image)  # gray
            image = self.train_trans0(image)  # horizen
            image = self.train_trans1(image)  # rotation
            image = self.train_trans2(image)  # sharp
            image = self.train_trans5(image)  # perspect
            
            image = preprocess_input(np.array(image, dtype='float32'))
            image = np.transpose(image, [2, 0, 1])
            image = torch.from_numpy(image)
            image = self.norm(image)
            
            image = self.train_trans3(image)
        
        elif self.mode == 'test':
            image = self.resize(image)
            image = preprocess_input(np.array(image, dtype='float32'))
            image = np.transpose(image, [2, 0, 1])
            image = torch.from_numpy(image)
            image = self.norm(image)
        
        return image, label

#     def __getitem__(self, index):
#         image = Image.open(self.paths[index])
#         label = self.labels[index]
#         image = self.crop(image)  # 160 160
#         # image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image = True)
#         if self.mode == 'train':
#             prob = np.random.rand()
#             if prob > 0.5:
#                 image = self.train_trans4(image)  # 90 90
#             image = self.resize(image)  # 160 160
#             image = self.train_trans0(image)  # horizen
#             image = self.train_trans1(image)  # rotation
#             image = self.train_trans2(image)  # sharp
#             image = self.train_trans5(image)  # perspect
            
#             image = np.array(image, dtype='float32')
#             image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125
#             # image = preprocess_input(np.array(image, dtype='float32'))
#             # image = np.transpose(image, [2, 0, 1])
#             image = torch.from_numpy(image)
#             # image = self.norm(image)
            
#             image = self.train_trans3(image)
        
#         elif self.mode == 'test':
#             image = self.resize(image)
#             image = np.array(image, dtype='float32')
#             image = (image.transpose((2, 0, 1)) - 127.5) * 0.0078125
#             # image = preprocess_input(np.array(image, dtype='float32'))
#             # image = np.transpose(image, [2, 0, 1])
#             image = torch.from_numpy(image)
#             # image = self.norm(image)
        
        return image, label
    
    def load_dataset(self):
        for path in self.lines:
            path_split = path.split(";")
            self.paths.append(path_split[1].split()[0])
            self.labels.append(int(path_split[0]))
        self.paths  = np.array(self.paths,dtype=object)
        self.labels = np.array(self.labels,dtype=object)
